// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/vector/constrained_vector_generator.h>

#include <gtest/gtest.h>

#include <pollux/common/constrained_generators.h>
#include <pollux/testing/vector/vector_test_base.h>

namespace kumo::pollux::fuzzer::test {

class ConstrainedVectorGeneratorTest : public testing::Test,
                                       public pollux::VectorBuilder {
 protected:
  static void SetUpTestCase() {
    memory::MemoryManager::testingSetInstance({});
  }

  template <TypeKind KIND>
  void testGenerateVectorsPrimitive(
      const TypePtr& type,
      const variant& excludedValue) {
    using T = typename TypeTraits<KIND>::NativeType;
    const uint32_t kSize = 1000;
    AbstractInputGeneratorPtr generator =
        std::make_shared<NotEqualConstrainedGenerator>(
            0,
            type,
            excludedValue,
            std::make_unique<RandomInputGenerator<T>>(0, type, 0.5));
    auto vector =
        ConstrainedVectorGenerator::generateConstant(generator, kSize, pool());
    EXPECT_EQ(vector->size(), kSize);
    EXPECT_EQ(vector->type_kind(), KIND);
    EXPECT_TRUE(vector->is_constant_encoding());
    EXPECT_TRUE(
        vector->is_null_at(0) ||
        vector->as<ConstantVector<T>>()->value_at(0) != excludedValue);

    vector = ConstrainedVectorGenerator::generateFlat(generator, kSize, pool());
    EXPECT_EQ(vector->size(), kSize);
    EXPECT_EQ(vector->type_kind(), KIND);
    EXPECT_TRUE(vector->is_flat_encoding());
    bool hasNull = false;
    for (auto i = 0; i < kSize; ++i) {
      if (vector->is_null_at(i)) {
        hasNull = true;
      } else {
        EXPECT_NE(vector->as<FlatVector<T>>()->value_at(i), excludedValue);
      }
    }
    EXPECT_TRUE(hasNull);
  }

  template <TypeKind KIND>
  void testGenerateVectorsComplex(const TypePtr& type) {
    using T = typename TypeTraits<KIND>::ImplType;
    const uint32_t kSize = 1000;
    AbstractInputGeneratorPtr generator =
        std::make_shared<RandomInputGenerator<T>>(0, type, 0.5);
    auto vector =
        ConstrainedVectorGenerator::generateFlat(generator, kSize, pool());
    EXPECT_EQ(vector->size(), kSize);
    EXPECT_EQ(vector->type(), type);
  }
};

TEST_F(ConstrainedVectorGeneratorTest, generateVectors) {
  testGenerateVectorsPrimitive<TypeKind::BIGINT>(BIGINT(), variant(0));
  testGenerateVectorsPrimitive<TypeKind::VARCHAR>(VARCHAR(), variant(""));

  testGenerateVectorsComplex<TypeKind::ARRAY>(
      ARRAY(ROW({MAP(VARCHAR(), BIGINT())})));
  testGenerateVectorsComplex<TypeKind::MAP>(
      MAP(ARRAY(BIGINT()), ROW({VARCHAR()})));
}

} // namespace kumo::pollux::fuzzer::test
